#!/usr/bin/env python
#coding=UTF-8
'''
Created on 2010-11-8

@author: zarra
'''

from common import Question,matrix2image,image2matrix,amplitude_plot
from numpy import pi,sin,exp,sqrt,abs,fromfunction
from numpy.fft import  fft2,ifft2
from q3_spatial_filtering.noise import Gaussian

def blur_gen(a,b,T):
    def h(u,v):
        p=(u*a+v*b)*pi
        #p[0,0] =1
        if p == 0: return T
        return T/p * sin(p)*exp(-1j*p)
    return h

class Q5(Question):
    def __init__(self,image,m='show'):
        super(self.__class__,self).__init__('Q5',image,m)

    def inverse_restoration(self,src,H,d): 
        ''' 逆滤波函数 src为输入图像的频率响应矩阵,H为原滤波器,d为对H的控制阈值'''
        src = src.copy()
        height,width = src.shape
        h_width = width/2
        h_height = height/2
        
       
        for i in range(height):
            for j in range(width):
                x=j-h_width
                y=i-h_height
                h=H(x,y)
                if  abs(h) > d:
                    src [i,j] /= h
                                         
        return src
    
    def inverse_image(self,image,H,d=0.0075):
        ''' 逆滤波函数image为输入图像,H为原滤波器,d为对H的控制阈值 '''
        fft = fft2(image2matrix(image))
        ivr = self.inverse_restoration(fft, H,d)
        
        fft_amp_img = amplitude_plot(ivr)
        
        ivr_img = matrix2image(ifft2(ivr))
        
        return ivr_img,fft_amp_img
        
    
    def wiener(self,src,H,K):
        ''' 维纳滤波函数 src为输入图像的频率响应矩阵,H为原滤波器,K为控制参数  返回所得的频率响应矩阵'''
        src=src.copy()
        height,width = src.shape
        h_width = width/2
        h_height = height/2
        for i in range(height):
            for j in range(width):
                x=j-h_width
                y=i-h_height
                h=H(x,y)
                if  abs(h) > K:
                    abs_h2 = abs(h)
                    abs_h2 *= abs_h2
                    src [i,j] *=  abs_h2/(h*(abs_h2+K))
                else:
                    src [i,j] = 0
        return src
    
    def wiener_image(self,m,H,K):
        ''' 维纳滤波函数 src为输入图像的频率响应矩阵,H为原滤波器,K为控制参数  返回所得的可输出图像'''
        wiener = self.wiener(m, H, K)
        wiener_img = matrix2image(ifft2(wiener))
        return wiener_img
            
    def blur(self,src,H):
        height,width = src.shape
        h_width = width/2
        h_height = height/2
#        def f(x,y):
#            print x
#            print '----------'
#            print y
#            print '----------'
#            p = pi*(0.1*x+0.1*y)
#            return 1/p * sin(p)*exp(-1j*p)
#        m_h=fromfunction(f, (height,width),dtype=complex)
#        print m_h
        for i in range(height):
            for j in range(width):
                h=H(j-h_width,i-h_height)
                src [i,j] *= h 
                              
    def action(self):
        image = self.image

        self.display(image, 'orig')
        H= blur_gen(0.1, 0.1, 1.0) #模糊函数
        
        gaus = Gaussian(0,sqrt(650)) #高斯噪声发生器
        gaus2 = Gaussian(0,sqrt(4))
        
        fft = fft2(image2matrix(image))
        
        fft_amp_img = amplitude_plot(fft)#原函数功率谱
        self.display(fft_amp_img, 'fft')
        
        self.blur(fft, H) #进行模糊

        ifft = ifft2(fft)
        blur_image = matrix2image(ifft)
        self.display(blur_image, 'blur')
        
        i,f = self.inverse_image(blur_image, H) #逆滤波
        
        self.display(i, 'inverse_for_blur')
        self.display(f, 'fft_of_inverse_for_blur')
    
        n_gaus = gaus(blur_image) #产生高斯噪声
        n_gaus2 = gaus2(blur_image)
        self.display(n_gaus, 'blur_gaus_noise_650')
        self.display(n_gaus2, 'blur_gaus_noise_4')
        i,f = self.inverse_image(n_gaus, H) #对带有噪声的图像再次逆滤波
        self.display(i, 'inverse_for_blur_noise_650')
        self.display(f, 'fft_of_inverse_for_blur_noise_650')
        
        i,f = self.inverse_image(n_gaus2, H)
        self.display(i, 'inverse_for_blur_noise_4')
        self.display(f, 'fft_of_inverse_for_blur_noise_4')
        fft_n_gaus = fft2(image2matrix(n_gaus))
        fft_n_gaus2 = fft2(image2matrix(n_gaus2))
        
        #维纳滤波
        K=0.01
        i=self.wiener_image(fft_n_gaus,H,K)   
        self.display(i, 'wiener_restoration_for_blur_noise_650_K=%.04f'%K)
        i=self.wiener_image(fft_n_gaus2,H,K) 
        self.display(i, 'wiener_restoration_for_blur_noise_4_K=%.04f'%K)
        K=0.005
        i=self.wiener_image(fft_n_gaus,H,K)   
        self.display(i, 'wiener_restoration_for_blur_noise_650_K=%.04f'%K)
        i=self.wiener_image(fft_n_gaus2,H,K) 
        self.display(i, 'wiener_restoration_for_blur_noise_4_K=%.04f'%K)
        K=0.001
        i=self.wiener_image(fft_n_gaus,H,K)   
        self.display(i, 'wiener_restoration_for_blur_noise_650_K=%.04f'%K)
        i=self.wiener_image(fft_n_gaus2,H,K) 
        self.display(i, 'wiener_restoration_for_blur_noise_4_K=%.04f'%K)


        
if __name__ == '__main__':
    q = Q5('fig6.jpg','save')
    q()
   
